{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a28e02f-ace3-41db-a115-3ae58a9b6d4d",
   "metadata": {},
   "source": [
    "## Hierarchical Agents with LlamaIndex\n",
    "\n",
    "The goal of this notebook is to demonstrate and compare the use of hierachical agents with LlamaIndex. We will set up a few tools and compare the capabilities of the agent when using the resources it has available for retrieval as agents vs tools.\n",
    "\n",
    "### Setup OpenAI Agent\n",
    "To start we will import the OpenAI agent we will be using across examples, and authenticate with our key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "50d77e7f-762c-497d-8a1f-bf5826e85dba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up OpenAI\n",
    "import openai\n",
    "from llama_index.agent import OpenAIAgent\n",
    "\n",
    "openai.api_key = \"sk-your-key\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b55495f1-0ea6-4d6f-879e-ecd393a0d6b0",
   "metadata": {},
   "source": [
    "### Setup Database Tool\n",
    "\n",
    "The first data source we will be setting up as a tool is the Database tool. This tool connects to a SQL database and can provide information about the tables, and run SQL queries against the database. One of the tools, `describe tables`, lists all of the tables and their columns, as seen below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c4c3dd4-5fff-4d97-ace6-e9f850c64b00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CREATE TABLE airports (\n",
      "\t\"Orig\" TEXT, \n",
      "\t\"Name\" TEXT, \n",
      "\t\"AirportCountry\" TEXT, \n",
      "\t\"Airport1Latitude\" TEXT, \n",
      "\t\"Airport1Longitude\" TEXT, \n",
      "\t\"TotalTravelers\" INTEGER\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "CREATE TABLE gdp (\n",
      "\t\"CountryName\" TEXT, \n",
      "\t\"MillionsGDP\" INTEGER\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "CREATE TABLE population (\n",
      "\t\"Country\" TEXT, \n",
      "\t\"ThousandsPopulation\" INTEGER\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "HEA, Herat, Afghanistan, 34.20690155, 62.22669983, 22041\n"
     ]
    }
   ],
   "source": [
    "# Import and initialize our tool spec\n",
    "from llama_hub.tools.database.base import DatabaseToolSpec\n",
    "from sqlalchemy import create_engine\n",
    "\n",
    "db_tool_spec = DatabaseToolSpec(engine=create_engine(\"sqlite:///countries.db\"))\n",
    "print(db_tool_spec.describe_tables())\n",
    "print(db_tool_spec.load_data(\"SELECT * from airports LIMIT 1\")[0].text)\n",
    "database_tools = db_tool_spec.to_tool_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc4559cd-6576-4e1a-a65d-35a35c91bed9",
   "metadata": {},
   "source": [
    "For the purposes of this demo, we are using a database of airports, GDP and population from the [world bank](https://datacatalog.worldbank.org/home). The airports table contains international airports, with the country they are located in and the number of travelers that passed through in 2019. The population and GDP tables contain population and GDP data of each country.\n",
    "\n",
    "In addition to the airports database, we will be using the wikipedia tool.\n",
    "\n",
    "### Wikipedia Query Engine Tool\n",
    "\n",
    "The wikipedia tool allows our agent to search wikipedia and pull entire articles based on the search. This is very useful for general knowledge retrieval, however most Wikipedia articles are nearly too large or too large to fit into context windows, and if our agent is retrieving multiple articles we quickly run out of tokens.\n",
    "\n",
    "Therefore, we combine two tools to give our Agent access to Wikipedia. We take the wikipedia `search_data` tool and wrap it with the `LoadAndSearchTool`. This results in two tools, where the first loads the wikipedia pages into a vector store, and the second acts as a query engine to retrieve data from the vector store. We will set this up and test it out below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a806dd33-f2ea-4187-9782-81154b3b6c53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Content loaded! You can now search the information using read_search_data\n",
      "The Denver Nuggets won the NBA 2023 playoffs.\n"
     ]
    }
   ],
   "source": [
    "# Import and initialize our tool spec\n",
    "from llama_hub.tools.wikipedia.base import WikipediaToolSpec\n",
    "from llama_index.tools.tool_spec.load_and_search.base import LoadAndSearchToolSpec\n",
    "\n",
    "wiki_spec = WikipediaToolSpec()\n",
    "# Get the search wikipedia tool\n",
    "search_wikipedia = wiki_spec.to_tool_list()[1]\n",
    "# Create the Agent with our tools\n",
    "wikipedia_tool_spec = LoadAndSearchToolSpec.from_defaults(search_wikipedia)\n",
    "\n",
    "print(wikipedia_tool_spec.load(\"NBA 2023 playoffs\"))\n",
    "print(wikipedia_tool_spec.read(\"Who won the NBA 2023 playoffs\"))\n",
    "wikipedia_tools = wikipedia_tool_spec.to_tool_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0c2753c-7e7b-4064-bfd8-7888fe9bf932",
   "metadata": {},
   "source": [
    "### Flat Agent\n",
    "\n",
    "We will now create two agents, one that directly uses the database tool, and another that has a dedicated agent for using the database tool, and compare the effectiveness.\n",
    "\n",
    "The \"flat\" is relatively simple to set up, we just pass the tools to our OpenAI Agent abstraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b8ce9049-c6ae-4ec9-908e-99ae772b4bc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "flat_agent = OpenAIAgent.from_tools(\n",
    "    [*wikipedia_tools, *database_tools],\n",
    "    system_prompt=\"\"\"\n",
    "        You are a specialized agent with access to wikipedia and a database containing information on\n",
    "        populations, gdp and the airports of countries. \n",
    "        When asked questions you should always attempt to use the database to gather information, and\n",
    "        wikipedia as a backup when the database does not contain the desired information.\n",
    "        Break any questions down into steps so that you can answer using a mix of the database and wikipedia.\n",
    "    \"\"\",\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bca4c627-dfd9-4eff-9972-82778ec98c61",
   "metadata": {},
   "source": [
    "### Creating the Hierarchical Agent\n",
    "\n",
    "Creating the hierarchical agent is somewhat more involved. One of the advantages of hierarchical agents is we can provide more specfic prompts to the sub-agents and for the agents as tools. For each sub-agent in a hierarchical agent we want to make, we make an agent with the tools it should have available and give the sub-agent a specfic system prompt for how it can use the tools. Then we wrap the sub-agent with the QueryEngineTool abstraction so that it can be passed as a tool to the higher level agent.\n",
    "\n",
    "We repeat this process to create a database and wikipedia agent tool that are passed to the hierarchical agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4825d3f7-1349-4621-97a2-cf994cc79564",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.query_engine import QueryEngineTool\n",
    "\n",
    "# Recreate the wikipedia tool spec so that our agents have seperate vector stores.\n",
    "wiki_spec = WikipediaToolSpec()\n",
    "search_wikipedia = wiki_spec.to_tool_list()[1]\n",
    "wikipedia_tool_spec = LoadAndSearchToolSpec.from_defaults(search_wikipedia)\n",
    "\n",
    "\n",
    "# Create the Database Agent\n",
    "database_agent = OpenAIAgent.from_tools(\n",
    "    [*database_tools],\n",
    "    system_prompt=f\"\"\"\n",
    "        You are a specialized agent designed to assist users query a SQL database.\n",
    "        \n",
    "        Here is a complete description of the SQL database you have access to:\n",
    "        \n",
    "        {db_tool_spec.describe_tables()}\n",
    "    \"\"\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "# Convert the Database Agent to a Tool\n",
    "database_agent_tool = QueryEngineTool.from_defaults(\n",
    "    database_agent,\n",
    "    name=\"database_agent\",\n",
    "    description=\"\"\"\n",
    "        This agent can answer questions about airport traffic, country GDP and population using a database\n",
    "    \"\"\",\n",
    ")\n",
    "\n",
    "# Create the Wikipedia Agent\n",
    "wikipedia_agent = OpenAIAgent.from_tools(\n",
    "    [*wikipedia_tools],\n",
    "    system_prompt=\"\"\"\n",
    "        You are a specialized agent designed to search for wikipedia pages and then \n",
    "        query the retrieved articles to answer questions passed to you by the user.\n",
    "        \n",
    "        This agent should be passed a specfic question to be answered using the wikipedia knowledge sources\n",
    "    \"\"\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "# Convert the Wikipedia Agent to a Tool\n",
    "wikipedia_agent_tool = QueryEngineTool.from_defaults(\n",
    "    wikipedia_agent,\n",
    "    name=\"wikipedia_agent\",\n",
    "    description=f\"\"\"\n",
    "        This agent can search wikipedia and parse articles to answer any questions on your behalf\n",
    "    \"\"\",\n",
    ")\n",
    "\n",
    "# Create the Hierarchical agent with the database agent\n",
    "hierarchical_agent = OpenAIAgent.from_tools(\n",
    "    [wikipedia_agent_tool, database_agent_tool],\n",
    "    system_prompt=\"\"\"\n",
    "        You are a specialized agent with access to wikipedia and a database containing information on\n",
    "        populations, gdp and the airports of countries. \n",
    "        When asked questions you should always attempt to use the database to gather information, and\n",
    "        wikipedia as a backup when the database does not contain the desired information.\n",
    "        Break any questions down into steps so that you can answer using a mix of the database and wikipedia.\n",
    "    \"\"\",\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0765ba37-51fd-4471-9333-f4d010a613c1",
   "metadata": {},
   "source": [
    "### Comparing the Agents\n",
    "\n",
    "One of the first issues the flat agent will run into is it doesn't have knowledge of the SQL schema. An advantage of the hierarchical agent is that we were able to supply the SQL schema in the system prompt for the database_agent. Because of this, the flat agent will struggle to use the database until we have it call the `describe_tables` tool to see the schema:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "34a620c8-1b3c-43f5-b08e-f684fef11f1b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: database_agent with args: {\n",
      "  \"input\": \"country with third highest gdp\"\n",
      "}\n",
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT CountryName, MillionsGDP FROM gdp ORDER BY MillionsGDP DESC LIMIT 3\"\n",
      "}\n",
      "Got output: [Document(id_='dfd1ced4-9c6c-478d-9999-fb7c95dd23eb', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='d85186287ae3c721bb63e6bacf6442aa9f3355d9a68ddb936ebcf1a6f50657d0', text='United States, 25462700', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='0b4b3b36-b9b1-4d83-9f9f-9000d95c95d5', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='52ab443d5c85d54558d1e749735171a48c35ef431492df0867fe8ab0dc0f4214', text='China, 17963171', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='f46732e5-b157-4213-b8ea-42f833171c66', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='3b545024dbd87ec2fc6992eaa8301da8103e81713af7dc3f3464f28c6e82b7e7', text='Japan, 4231141', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "Got output: The country with the third highest GDP is Japan, with a GDP of 4,231,141 million.\n",
      "========================\n",
      "The country with the third highest GDP is Japan.\n"
     ]
    }
   ],
   "source": [
    "print(hierarchical_agent.chat(\"what country has the third highest gdp\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "74b688a1-68b4-457d-93a0-40b41eed7e83",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT country_name, gdp FROM countries ORDER BY gdp DESC LIMIT 3\"\n",
      "}\n"
     ]
    },
    {
     "ename": "OperationalError",
     "evalue": "(sqlite3.OperationalError) no such table: countries\n[SQL: SELECT country_name, gdp FROM countries ORDER BY gdp DESC LIMIT 3]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "<clipped error>",
      "\u001b[0;31mOperationalError\u001b[0m: (sqlite3.OperationalError) no such table: countries\n[SQL: SELECT country_name, gdp FROM countries ORDER BY gdp DESC LIMIT 3]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)"
     ]
    }
   ],
   "source": [
    "print(flat_agent.chat(\"what country has the third highest gdp\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "818f37cb-76a3-48ce-8f81-2a831d48ad3f",
   "metadata": {},
   "source": [
    "### Hierachical Agent Advantage: Additional specfic context\n",
    "\n",
    "The flat agent was unable to properly form the SQL query as it hadn't yet retrieved the table schema. In comparison, we provided the database agent the database schema as part of it's system prompt, as it's main focus is interacting with the database. We can ask the flat agent to load the table schema and ask the question again to get the correct answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b4cd6e4e-33df-4cc2-b377-9a74e192d29a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: describe_tables with args: {}\n",
      "Got output: \n",
      "CREATE TABLE airports (\n",
      "\t\"Orig\" TEXT, \n",
      "\t\"Name\" TEXT, \n",
      "\t\"AirportCountry\" TEXT, \n",
      "\t\"Airport1Latitude\" TEXT, \n",
      "\t\"Airport1Longitude\" TEXT, \n",
      "\t\"TotalTravelers\" INTEGER\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "CREATE TABLE gdp (\n",
      "\t\"CountryName\" TEXT, \n",
      "\t\"MillionsGDP\" INTEGER\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "CREATE TABLE population (\n",
      "\t\"Country\" TEXT, \n",
      "\t\"ThousandsPopulation\" INTEGER\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "========================\n",
      "There are three available tables in the database:\n",
      "\n",
      "1. `airports` table:\n",
      "   - Columns: \"Orig\" (origin), \"Name\", \"AirportCountry\", \"Airport1Latitude\", \"Airport1Longitude\", \"TotalTravelers\"\n",
      "   - This table contains information about airports, including their names, countries, coordinates, and total number of travelers.\n",
      "\n",
      "2. `gdp` table:\n",
      "   - Columns: \"CountryName\", \"MillionsGDP\"\n",
      "   - This table contains information about the GDP (Gross Domestic Product) of countries, including the country name and the GDP in millions.\n",
      "\n",
      "3. `population` table:\n",
      "   - Columns: \"Country\", \"ThousandsPopulation\"\n",
      "   - This table contains information about the population of countries, including the country name and the population in thousands.\n",
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT CountryName FROM gdp ORDER BY MillionsGDP DESC LIMIT 3\"\n",
      "}\n",
      "Got output: [Document(id_='bf40672d-5d1d-4fde-9dcf-edfdbb3c8e97', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='824aba9f20b7a8581946ff24092b95c8728747192b25d38eae0185cec4b66057', text='United States', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='c6c35f45-0e3b-41db-aa9f-e05728b1268e', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='99d48fd32be248a61bf8694eee10f141a2ce750e9d5daad17ac71946b71a7a26', text='China', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='01c2dbe4-b2e5-42d1-a5f9-004e4c118de3', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='6bc95f23434f6d132ac882d56019ba633401b420c3052c47f79555eb00b0f717', text='Japan', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: read_search_data with args: {\n",
      "  \"query\": \"What is the GDP of Japan?\"\n",
      "}\n",
      "Got output: I'm sorry, but I don't have access to real-time information. My responses are based on the context provided and not on current data.\n",
      "========================\n",
      "The country with the third highest GDP is Japan.\n"
     ]
    }
   ],
   "source": [
    "print(flat_agent.chat(\"describe all of the available tables\"))\n",
    "print(flat_agent.chat(\"what country has the third highest gdp\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d128cc9-9481-4851-a300-2c9e1ea8919c",
   "metadata": {},
   "source": [
    "The flat agent was able to form the correct query after being explicity asked by the user, however it still made an extra query to the wikipedia tool when the data could have been retrieved from the database.\n",
    "\n",
    "### Comparing ability to write Join queries\n",
    "\n",
    "We ask the agents to execute a SQL query that requires a join:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "65d3c6c3-3abe-42c6-939c-a6fb2b2be765",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: database_agent with args: {\n",
      "  \"input\": \"country with highest gdp and population under 10 million\"\n",
      "}\n",
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT gdp.CountryName, gdp.MillionsGDP, population.ThousandsPopulation FROM gdp JOIN population ON gdp.CountryName = population.Country WHERE population.ThousandsPopulation < 10000 ORDER BY gdp.MillionsGDP DESC LIMIT 1\"\n",
      "}\n",
      "Got output: [Document(id_='726b0010-5f0b-4843-b38a-55e51d32fdd7', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='d176dff8313dab084435594a2f2c4071819c5ccf5b6abb2d03d4889d0ae4ef8e', text='Switzerland, 807706, 8770', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "Got output: The country with the highest GDP and a population under 10 million is Switzerland, with a GDP of 807,706 million and a population of 8,770 thousand.\n",
      "========================\n",
      "The country with under 10 million people and the highest GDP is Switzerland.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    hierarchical_agent.chat(\n",
    "        \"What country with under 10 million people has the highest gdp\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fc9a4f8c-e202-48ae-bb39-b83bd7e6fb65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT CountryName, MillionsGDP FROM gdp WHERE MillionsGDP IS NOT NULL AND ThousandsPopulation < 10000 ORDER BY MillionsGDP DESC LIMIT 1\"\n",
      "}\n"
     ]
    },
    {
     "ename": "OperationalError",
     "evalue": "(sqlite3.OperationalError) no such column: ThousandsPopulation\n[SQL: SELECT CountryName, MillionsGDP FROM gdp WHERE MillionsGDP IS NOT NULL AND ThousandsPopulation < 10000 ORDER BY MillionsGDP DESC LIMIT 1]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "<clipped error>",
      "\u001b[0;31mOperationalError\u001b[0m: (sqlite3.OperationalError) no such column: ThousandsPopulation\n[SQL: SELECT CountryName, MillionsGDP FROM gdp WHERE MillionsGDP IS NOT NULL AND ThousandsPopulation < 10000 ORDER BY MillionsGDP DESC LIMIT 1]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)"
     ]
    }
   ],
   "source": [
    "print(flat_agent.chat(\"What country with under 10 million people has the highest gdp\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52fe5f89-7268-4916-a115-3ad506b67bf7",
   "metadata": {},
   "source": [
    "### Hierarchical Agent Advantage: Better ability to craft queries\n",
    "\n",
    "The flat agent was unable to come up with the correct join query to combine the population and gdp tables, whereas the hierarchical agent executed the query correctly and obtained the right answer. One reason that the hierarchical agent performed better here may be that schema was closer to the question in the context window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "de5ec146-f783-4b75-9754-d8a33ec3d89e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: database_agent with args: {\n",
      "  \"input\": \"country with third highest population\"\n",
      "}\n",
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT Country, ThousandsPopulation FROM population ORDER BY ThousandsPopulation DESC LIMIT 3\"\n",
      "}\n",
      "Got output: [Document(id_='b3e294be-0cea-446e-b161-06d760c93f06', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='3eb098f7aa2d847b30fa9837d8798f637148c3c6a120ae12612114412660bfce', text='India, 1417173', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='e43f63ed-2a10-44af-9983-a37fd42c437b', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='72c2fca9e61ac0cbbfb69f533b765fc241006300363ad524bc140aa9f903b807', text='China, 1412175', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='46487305-0665-4bba-bfa7-f6d7bad692cd', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='f13a50923bd5a0240b6675774069830540ed10d4dbffd34df0ca4481786cd7ad', text='United States, 333288', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "Got output: The country with the third highest population is the United States, with a population of 333,288 thousand.\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: wikipedia_agent with args: {\n",
      "  \"input\": \"current leader of the United States\"\n",
      "}\n",
      "=== Calling Function ===\n",
      "Calling function: search_data with args: {\n",
      "  \"query\": \"President of the United States\"\n",
      "}\n",
      "Got output: Content loaded! You can now search the information using read_search_data\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: read_search_data with args: {\n",
      "  \"query\": \"Who is the current leader of the United States?\"\n",
      "}\n",
      "Got output: Joe Biden is the current leader of the United States.\n",
      "========================\n",
      "Got output: The current leader of the United States is Joe Biden.\n",
      "========================\n",
      "The current leader of the country with the third highest population, which is the United States, is Joe Biden.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    hierarchical_agent.chat(\n",
    "        \"Who is the current leader of the country with the third highest population\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "56a1cbe3-31a6-4b0d-9844-ce8ea494368b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT Country FROM population ORDER BY ThousandsPopulation DESC LIMIT 3\"\n",
      "}\n",
      "Got output: [Document(id_='af84a01d-816f-45ed-bc14-b8ba61b82455', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='abbc17675c6627a122d3c0160f6cce46a9f841a6469eb47ab5d87a6499fd86b7', text='India', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='6a6948bb-c64d-4d07-906e-ba3b04291302', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='99d48fd32be248a61bf8694eee10f141a2ce750e9d5daad17ac71946b71a7a26', text='China', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='23f035b3-4663-4e67-b68d-c3f3d9ff15fc', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='824aba9f20b7a8581946ff24092b95c8728747192b25d38eae0185cec4b66057', text='United States', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: search_data with args: {\n",
      "  \"query\": \"current leader of India\"\n",
      "}\n",
      "Got output: Content loaded! You can now search the information using read_search_data\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: read_search_data with args: {\n",
      "  \"query\": \"Who is the current leader of India?\"\n",
      "}\n",
      "Got output: The current leader of India is the prime minister.\n",
      "========================\n",
      "The country with the third highest population is India. The current leader of India is the Prime Minister.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    flat_agent.chat(\n",
    "        \"Who is the current leader of the country with the third highest population\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f54d92c1-ceae-4562-819e-e1725321b8ea",
   "metadata": {},
   "source": [
    "### Hierarchical Agent Advantage: Ability to interpret results\n",
    "\n",
    "The hierarchical agent correctly took the last value in the query as the third highest population, where as the flat agent took the first result which is actually the highest population. It appears that they were both able to use the wikipedia tool equally well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2054f017-3592-403b-98c9-b42104980c47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: database_agent with args: {\n",
      "  \"input\": \"fifth busiest airport in Spain\"\n",
      "}\n",
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT Name, TotalTravelers FROM airports WHERE AirportCountry = 'Spain' ORDER BY TotalTravelers DESC LIMIT 5\"\n",
      "}\n",
      "Got output: [Document(id_='a1d7c549-4fa0-4ea8-8e6f-e536bdf692ee', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='d3633a93029fc54b48f25e444778f0dbe985b57f76f36e8e79e6e734b14aee90', text='Adolfo Suarez-Barajas, 26811681', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='7a638f65-7dc9-432b-80ea-6db76e2c9d46', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='cec3baa9dbbb70249033fb6c6995520286d54b3144e50e33db45f3f2bf4b87bb', text='Barcelona, 23189184', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='bc4a8fe2-6909-4336-a4c0-14c44ddd4fd0', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='90047eb6de8f4de7e3f7317c1863e53bfc3a15330dd79c731c8e8920309959fd', text='Palma de Mallorca, 11811792', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='aec768e4-009b-4cc7-ad67-9769711f3b42', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='b653e7d4f830a6fb8dd2537883524e37c6e4e6a557470897b87ce1b27e8d6550', text='Malaga Airport, 9401882', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='8996906b-d7b7-4269-81b2-e6ed505c024b', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='81d7a3092dddf245c721f2d12ce72729c9a0cfcdfa8f247dd6a10b8305209fd6', text='Alicante Airport, 7500704', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "Got output: The fifth busiest airport in Spain is Alicante Airport, with a total of 7,500,704 travelers.\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: wikipedia_agent with args: {\n",
      "  \"input\": \"Alicante Airport\"\n",
      "}\n",
      "=== Calling Function ===\n",
      "Calling function: search_data with args: {\n",
      "  \"query\": \"Alicante Airport\"\n",
      "}\n",
      "Got output: Content loaded! You can now search the information using read_search_data\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: read_search_data with args: {\n",
      "  \"query\": \"What is Alicante Airport?\"\n",
      "}\n",
      "Got output: Alicante Airport, officially known as Alicante–Elche Miguel Hernández Airport, is an international airport located in Spain. It is situated about 9 km southwest of the city of Alicante and about 10 km east of the city of Elche. The airport serves as a major transportation hub in southeastern Spain, catering to both the southern part of the Valencian Community and the Region of Murcia. It is one of the busiest airports in Europe, with a significant portion of its passenger traffic being international. The airport has a single terminal, Terminal N, which was opened in 2011 and has a capacity of 20 million passengers per year. Alicante Airport is served by several airlines, including Air Nostrum, Norwegian Air Shuttle, Ryanair, and Vueling.\n",
      "========================\n",
      "Got output: Alicante Airport, officially known as Alicante–Elche Miguel Hernández Airport, is an international airport located in Spain. It is situated about 9 km southwest of the city of Alicante and about 10 km east of the city of Elche. The airport serves as a major transportation hub in southeastern Spain, catering to both the southern part of the Valencian Community and the Region of Murcia. It is one of the busiest airports in Europe, with a significant portion of its passenger traffic being international. The airport has a single terminal, Terminal N, which was opened in 2011 and has a capacity of 20 million passengers per year. Alicante Airport is served by several airlines, including Air Nostrum, Norwegian Air Shuttle, Ryanair, and Vueling.\n",
      "========================\n",
      "The fifth busiest airport in Spain, Alicante Airport, was opened in 2011.\n"
     ]
    }
   ],
   "source": [
    "print(hierarchical_agent.chat(\"when was the fifth busiest airport in Spain built?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "36b03b3f-b13d-4b37-8346-f14e91c2173a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT Name FROM airports WHERE AirportCountry = 'Spain' ORDER BY TotalTravelers DESC LIMIT 5\"\n",
      "}\n",
      "Got output: [Document(id_='1c3bc3aa-9fc1-42b6-b97a-18e7504b497e', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='6e5f8e00b3cfc65c4147a4a548ed00dba6ee4756bb29604b58ed5fbb707d3fe1', text='Adolfo Suarez-Barajas', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='5e7c5f25-e2eb-47be-bd82-6fe675df1b70', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='b7ab1ee21e956cea7f8214063fc92c17c8be117642ba67a19f7d1fe4d7736080', text='Barcelona', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='6268cbf3-3226-44c0-8412-1b2dfaa60d09', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='2cf83441259ffc3fb9681e6c41504271a0b3e2174e265d6ed94ccd7e921e03ce', text='Palma de Mallorca', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='4f30e752-951a-4194-a345-957a758edd56', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='68a5b7d0ba86be2729f9221653f2b14d4af8987d639329527be573e769725be3', text='Malaga Airport', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='a222ed51-3c47-45d4-9c4b-263d12e37bb1', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='b531246331b1458caff6ae2d9ba8cc7636aa384e1d6fafa846d1e96fcc083df5', text='Alicante Airport', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: search_data with args: {\n",
      "  \"query\": \"When was Adolfo Suarez-Barajas airport built?\"\n",
      "}\n",
      "Got output: Content loaded! You can now search the information using read_search_data\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: read_search_data with args: {\n",
      "  \"query\": \"When was Adolfo Suarez-Barajas airport built?\"\n",
      "}\n",
      "Got output: Adolfo Suárez Madrid–Barajas Airport was built in 1927.\n",
      "========================\n",
      "The fifth busiest airport in Spain is Adolfo Suarez-Barajas Airport. It was built in 1927.\n"
     ]
    }
   ],
   "source": [
    "print(flat_agent.chat(\"when was the fifth busiest airport in Spain built?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0bed80d-f69e-4fbf-a532-4f195f91887d",
   "metadata": {},
   "source": [
    "### No agent advantage\n",
    "\n",
    "Both agents were unable to properly answer this question, as the flat agent again took the first result instead of the fifth. The hierarchical agent selected the correct airport, however it was not able to use the wikipedia tool effecitvely to determine the date of construction of the airport, instead providing the date an additional terminal was constructed. It's possible with more prompt tuning we could get the hierarchical agent to use the wikipedia sub-agent more effectively."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
